# 主函数
import pickle
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from stock import Bar
from strategy import Strategy
from backtest import Backtest
from visualization import Visualization

import math
from pylab import mpl   # 从pylab导入子模块mpl
mpl.rcParams['font.sans-serif']=['KaiTi']    # 以楷体的字体显示中文
mpl.rcParams['axes.unicode_minus']=False # 解决保存图像是负号'-'显示为方块的问题

def getSharpRatio(codelist, all_data, rate, startDateTime, endDateTime):
    """
      根据股票的夏普比例进行选股，用于策略
      :param codelist: list, 股票代码
      :param all_data: dict, 初始所有数据
      :param rate: float, 对标的无风险利率
      :param startDateTime: str, 用来计算夏普比率的初始时间
      :param endDateTime: str, 用来计算夏普比率的结束时间
      :return: DataFrame，index是股票代码，columns是夏普比率，内容是每只股票的夏普比率
      """
    sharp = {}
    for code in codelist:
        stock = Bar(code,["change_pct"])
        stock.selectStockInformation(all_data, startDateTime, endDateTime)
        data = stock.stock_data["change_pct"].dropna()
        day_num = len(data)
        sharp[code] = (data.mean()*day_num-rate)/(data.std()*math.sqrt(day_num))
        sharpdf = pd.DataFrame(sharp, index=["sharp"]).T
        sharpdf = sharpdf.sort_values("sharp", ascending=False)  # 将sharpdf按夏普比率降序排序
    return sharpdf


# 针对策略2，寻找收益最高的参数
def getbest(Bars, cash, circle):
    best_return = -2
    best_up = 0
    best_down = 0
    for i in np.arange(0, 2, 0.1):
        for j in np.arange(0, 1, 0.1):
            print(i,j)
            s = Strategy()
            b = Backtest(Bars, cash, circle)
            df = b.get_calendar_data(s.trade2(Bars, up=i, down=j))
            allreturn = (df.iloc[-1, -1] - cash) / cash
            if allreturn > best_return:
                best_up = i
                best_down = j
                best_return = allreturn
                print(i, j, best_return)
    return best_up, best_down


if __name__ == "__main__":
    file_path = "./data.pkl"
    with open(file_path, "rb") as f:
        all_data = pickle.load(f)
    codelist = all_data["2010-01-04"].index.tolist()
    indicatorlist = ["open", "close", "high", "low", "volume", "change_pct"]

    startDateTime = "2019-07-01"
    endDateTime = "2019-12-31"

    stocklist = getSharpRatio(codelist, all_data, rate=0.03, startDateTime="2018-07-01", endDateTime="2018-12-31").index
    stockNum = 3

    Bars = []  # 创建多个类对象
    for code in stocklist[0:stockNum]:  # 选定排名前几的股票作为股票池
        stock = Bar(code, indicatorlist)  # 调用Bar类
        stock.selectStockInformation(all_data, "2019-07-01", "2019-12-31")
        Bars.append(stock)

    cash = 1000000  # 初始投入资金
    circle = 5   # 调仓周期

    # 策略1
    s1 = Strategy()
    weight1 = s1.trade1(Bars)
    b1 = Backtest(Bars, cash, circle)
    record_df1 = b1.get_calendar_data(weight1)

    # 策略2
    s2 = Strategy()
    # bestparam = getbest(Bars, cash, circle)
    # weight2 = s2.trade2(Bars, up=bestparam[0], down=bestparam[1])
    weight2 = s2.trade2(Bars, up=1.9, down=0)
    b2 = Backtest(Bars, cash, circle)
    record_df2 = b2.get_calendar_data(weight2)


    # 导入上证指数数据
    benchmark_data = pd.read_excel("./沪深300和上证指数2010-2020历史数据.xlsx", sheet_name="上证指数", header=0, index_col=0)
    graph = Visualization()

    benchmark_return =graph.sum_return_single(benchmark_data, record_df1.index[0], record_df1.index[-1],
                               title="上证指数累计收益率曲线图", return_BaseOn="close")

    # 策略的累计收益率
    yielddf = graph.get_yield(record_df2, cash)

    # 策略和上证指数的累计收益率比较图
    list1 = benchmark_return["所选股票当日累计收益率"].tolist()
    list2 = yielddf["策略的累计收益率"].tolist()
    df_compare = pd.DataFrame({"上证指数累计收益率": list1, "策略的累计收益率": list2}, index=yielddf.index.tolist())
    df_compare.plot(kind='line', title=u'累计收益率对比图', figsize=(9, 9))
    plt.show()
    plt.close()


